#! /usr/bin/env python
#
# pygrub - simple python-based bootloader for Xen
#
# Copyright 2005-2006 Red Hat, Inc.
# Jeremy Katz <katzj@redhat.com>
#
# This software may be freely redistributed under the terms of the GNU
# general public license.
#
# You should have received a copy of the GNU General Public License
# along with this program; If not, see <http://www.gnu.org/licenses/>.
#

from __future__ import print_function

import os, sys, string, struct, tempfile, re, traceback, stat, errno
import copy
import ctypes, ctypes.util
import logging
import platform
import resource
import subprocess

import curses, _curses, curses.textpad, curses.ascii
import getopt

import xenfsimage
import grub.GrubConf
import grub.LiloConf
import grub.ExtLinuxConf

PYGRUB_VER = 0.7
FS_READ_MAX = 1024 * 1024
SECTOR_SIZE = 512

# Unless provided through the env variable PYGRUB_MAX_FILE_SIZE_MB, then
# this is the maximum filesize allowed for files written by the depriv
# pygrub
LIMIT_FSIZE = 128 << 20

# Unless provided through the env variable PYGRUB_MAX_RAM_USAGE_MB, then
# this is the maximum amount of memory allowed to be used by the depriv
# pygrub.
LIMIT_AS = 2 * LIMIT_FSIZE

CLONE_NEWNS = 0x00020000 # mount namespace
CLONE_NEWNET = 0x40000000 # network namespace
CLONE_NEWIPC = 0x08000000 # IPC namespace

def unshare(flags):
    if not sys.platform.startswith("linux"):
        print("skip_unshare reason=not_linux platform=%s", sys.platform, file=sys.stderr)
        return

    libc = ctypes.CDLL(ctypes.util.find_library('c'), use_errno=True)
    unshare_prototype = ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_int, use_errno=True)
    unshare = unshare_prototype(('unshare', libc))

    if unshare(flags) < 0:
        raise OSError(ctypes.get_errno(), os.strerror(ctypes.get_errno()))

    # It's very typical for systemd to mount / with MS_SHARED. That means
    # any events in the new namespace get propagated back to the parent.
    #
    # Undo it so that every mount done in the NS stay confined within it.
    subprocess.check_output(["mount", "--make-rprivate", "/"])

def bind_mount(src, dst, options):
    open(dst, "a").close() # touch

    rc = subprocess.call(["mount", "--bind", "-o", options, src, dst])
    if rc != 0:
        raise RuntimeError("bad_mount: src=%s dst=%s opts=%s" %
                           (src, dst, options))

def downgrade_rlimits():
    # Wipe the authority to use unrequired resources
    resource.setrlimit(resource.RLIMIT_NPROC,    (0, 0))
    resource.setrlimit(resource.RLIMIT_CORE,     (0, 0))
    resource.setrlimit(resource.RLIMIT_MEMLOCK,  (0, 0))

    max_ram_usage = LIMIT_AS
    if "PYGRUB_MAX_RAM_USAGE_MB" in os.environ:
        max_ram_usage = int(os.environ["PYGRUB_MAX_RAM_USAGE_MB"]) << 20
    resource.setrlimit(resource.RLIMIT_AS,  (max_ram_usage, max_ram_usage))

    # py2's resource module doesn't know about resource.RLIMIT_MSGQUEUE
    #
    # TODO: Use resource.RLIMIT_MSGQUEUE after python2 is deprecated
    if sys.platform.startswith('linux'):
        RLIMIT_MSGQUEUE = 12
        resource.setrlimit(RLIMIT_MSGQUEUE, (0, 0))

    # The final look of the filesystem for this process is fully RO, but
    # note we have some file descriptor already open (notably, kernel and
    # ramdisk). In order to avoid a compromised pygrub from filling up the
    # filesystem we set RLIMIT_FSIZE to a high bound, so that the file
    # write permissions are bound.
    fsize = LIMIT_FSIZE
    if "PYGRUB_MAX_FILE_SIZE_MB" in os.environ:
        fsize = int(os.environ["PYGRUB_MAX_FILE_SIZE_MB"]) << 20

    resource.setrlimit(resource.RLIMIT_FSIZE, (fsize, fsize))

def depriv(output_directory, output, device, uid, path_kernel, path_ramdisk):
    # The only point of this call is to force the loading of libfsimage.
    # That way, we don't need to bind-mount it into the chroot
    rc = xenfsimage.init()
    if rc != 0:
        os.unlink(path_ramdisk)
        os.unlink(path_kernel)
        raise RuntimeError("bad_xenfsimage: rc=%d" % rc)

    # Create a temporary directory for the chroot
    chroot = tempfile.mkdtemp(prefix=str(uid)+'-', dir=output_directory) + '/'
    device_path = '/device'

    pid = os.fork()
    if pid:
        # parent
        _, rc = os.waitpid(pid, 0)

        for path in [path_kernel, path_ramdisk]:
            # If the child didn't write anything, just get rid of it,
            # otherwise we end up consuming a 0-size file when parsing
            # systems without a ramdisk that the ultimate caller of pygrub
            # may just be unaware of
            if rc != 0 or os.path.getsize(path) == 0:
                os.unlink(path)

        # Unshare(CLONE_NEWNS) ensures this is not required, but that's not
        # present on *BSD, so recursively unmount everything if needed.
        # Quietly.
        with open('/dev/null', 'w') as devnull:
            subprocess.call(["umount", "-f", chroot + device_path],
                            stdout=devnull, stderr=devnull)
            subprocess.call(["umount", "-f", chroot],
                            stdout=devnull, stderr=devnull)
        os.rmdir(chroot)

        sys.exit(rc)

    # By unsharing the namespace we're making sure it's all bulk-released
    # at the end, when the namespaces disappear. This means the kernel does
    # (almost) all the cleanup for us and the parent just has to remove the
    # temporary directory.
    unshare(CLONE_NEWNS | CLONE_NEWIPC | CLONE_NEWNET)

    # Set sensible limits using the setrlimit interface
    downgrade_rlimits()

    # We'll mount tmpfs on the chroot to ensure the deprivileged child
    # cannot affect the persistent state. It's RW now in order to
    # bind-mount the device, but note it's remounted RO after that.
    rc = subprocess.call(["mount", "-t", "tmpfs", "none", chroot])
    if rc != 0:
        raise RuntimeError("mount_tmpfs rc=%d dst=\"%s\"" % (rc, chroot))

    # Bind the untrusted device RO
    bind_mount(device, chroot + device_path, "ro,nosuid,noexec")

    rc = subprocess.call(["mount", "-t", "tmpfs", "-o", "remount,ro,nosuid,noexec,nodev", "none", chroot])
    if rc != 0:
        raise RuntimeError("remount_tmpfs rc=%d dst=\"%s\"" % (rc, chroot))

    # Drop superpowers!
    os.chroot(chroot)
    os.chdir('/')
    os.setgid(uid)
    os.setgroups([uid])
    os.setuid(uid)

    return device_path

def read_size_roundup(fd, size):
    if platform.system() != 'FreeBSD':
        return size
    st = os.fstat(fd)
    if (not stat.S_ISCHR(st.st_mode)):
        return size
    # Round up to sector size if it's a raw character device
    return (((size)+((SECTOR_SIZE)-1))&(~((SECTOR_SIZE)-1)))

def enable_cursor(ison):
    if ison:
        val = 2
    else:
        val = 0

    try:
        curses.curs_set(val)
    except _curses.error:
        pass

DISK_TYPE_RAW, DISK_TYPE_HYBRIDISO, DISK_TYPE_DOS = range(3)
def identify_disk_image(file):
    """Detect DOS partition table or HybridISO format."""
    fd = os.open(file, os.O_RDONLY)
    buf = os.read(fd, read_size_roundup(fd, 0x8006))
    os.close(fd)

    if len(buf) >= 512 and \
           struct.unpack("H", buf[0x1fe: 0x200]) == (0xaa55,):
        # HybridISO contains a DOS partition table for booting from USB devices, but really is an ISO image
        if len(buf) >= 0x8006 and buf[0x8001:0x8006] == 'CD001':
            return DISK_TYPE_HYBRIDISO
        return DISK_TYPE_DOS
    return DISK_TYPE_RAW

DK_LABEL_LOC=1
DKL_MAGIC=0xdabe
V_ROOT=0x2

def get_solaris_slice(file, offset):
    """Find the root slice in a Solaris VTOC."""

    fd = os.open(file, os.O_RDONLY)
    os.lseek(fd, offset + (DK_LABEL_LOC * SECTOR_SIZE), 0)
    buf = os.read(fd, 512)
    os.close(fd)
    if struct.unpack("<H", buf[508:510])[0] != DKL_MAGIC:
        raise RuntimeError("Invalid disklabel magic")

    nslices = struct.unpack("<H", buf[30:32])[0]

    for i in range(nslices):
        sliceoff = 72 + 12 * i
        slicetag = struct.unpack("<H", buf[sliceoff:sliceoff+2])[0]
        slicesect = struct.unpack("<L", buf[sliceoff+4:sliceoff+8])[0]
        if slicetag == V_ROOT:
            return slicesect * SECTOR_SIZE

    raise RuntimeError("No root slice found")

def get_fs_offset_gpt(file):
    fd = os.open(file, os.O_RDONLY)
    os.lseek(fd, SECTOR_SIZE, 0)
    buf = os.read(fd, 512)
    partcount = struct.unpack("<L", buf[80:84])[0]
    partsize = struct.unpack("<L", buf[84:88])[0]
    i = partcount
    offsets = []
    while i>0:
        buf = os.read(fd, read_size_roundup(fd, partsize))
        offsets.append(struct.unpack("<Q", buf[32:40])[0] * SECTOR_SIZE)
        i -= 1
    os.close(fd)
    return offsets

FDISK_PART_SOLARIS=0xbf
FDISK_PART_SOLARIS_OLD=0x82
FDISK_PART_GPT=0xee

def get_partition_offsets(file):
    image_type = identify_disk_image(file)
    if image_type == DISK_TYPE_RAW:
        # No MBR: assume whole disk filesystem, which is like a
        # single partition starting at 0
        return [0]
    elif image_type == DISK_TYPE_HYBRIDISO:
        # A HybridISO contains an ISO filesystem at 0 in addition
        # to the DOS partition table
        part_offs = [0]
    elif image_type == DISK_TYPE_DOS:
        part_offs = []
    else:
        raise ValueError('Unhandled image type returnd by identify_disk_image(): %d' % (image_type,))

    fd = os.open(file, os.O_RDONLY)
    buf = os.read(fd, 512)
    os.close(fd)
    for poff in (446, 462, 478, 494): # partition offsets

        # MBR contains a 16 byte descriptor per partition
        partbuf = buf[poff:poff+16]
        offset  = struct.unpack("<L", partbuf[8:12])[0] * SECTOR_SIZE
        type    = struct.unpack("<B", partbuf[4:5])[0]

        # offset == 0 implies this partition is not enabled
        if offset == 0:
            continue

        if type == FDISK_PART_SOLARIS or type == FDISK_PART_SOLARIS_OLD:
            try:
                offset += get_solaris_slice(file, offset)
            except RuntimeError:
                continue # no solaris magic at that offset, ignore partition

        if type == FDISK_PART_GPT:
            for offset in get_fs_offset_gpt(file):
                part_offs.append(offset)
            break

        # Active partition has 0x80 as the first byte.
        # If active, prepend to front of list, otherwise append to back.
        if struct.unpack("<c", buf[poff:poff+1]) == ('\x80',):
            part_offs.insert(0, offset)
        else:
            part_offs.append(offset)

    # We thought we had a DOS partition table, but didn't find any
    # actual valid partition entries. This can happen because an MBR
    # (e.g. grubs) may contain the same signature.
    if not part_offs: part_offs = [0]

    return part_offs

class GrubLineEditor(curses.textpad.Textbox):
    def __init__(self, screen, startx, starty, line = ""):
        screen.addstr(startx, starty, "> ")
        screen.noutrefresh()
        win = curses.newwin(1, 74, startx, starty + 2)
        curses.textpad.Textbox.__init__(self, win)

        self.line = list(line)
        self.pos = len(line)
        self.cancelled = False
        self.show_text()

    def show_text(self):
        """Show the text.  One of our advantages over standard textboxes
        is that we can handle lines longer than the window."""

        self.win.erase()
        p = self.pos
        off = 0
        while p > 70:
            p -= 55
            off += 55

        l = self.line[off:off+70]
        self.win.addstr(0, 0, string.join(l, ("")))
        if self.pos > 70:
            self.win.addch(0, 0, curses.ACS_LARROW)

        self.win.move(0, p)

    def do_command(self, ch):
        # we handle escape as well as moving the line around, so have
        # to override some of the default handling

        self.lastcmd = ch
        if ch == 27: # esc
            self.cancelled = True
            return 0
        elif curses.ascii.isprint(ch):
            self.line.insert(self.pos, chr(ch))
            self.pos += 1
        elif ch == curses.ascii.SOH:  # ^a
            self.pos = 0
        elif ch in (curses.ascii.STX,curses.KEY_LEFT):
            if self.pos > 0:
                self.pos -= 1
        elif ch in (curses.ascii.BS,curses.KEY_BACKSPACE):
            if self.pos > 0:
                self.pos -= 1
                if self.pos < len(self.line):
                    self.line.pop(self.pos)
        elif ch == curses.ascii.EOT:                           # ^d
            if self.pos < len(self.line):
                self.line.pop(self.pos)
        elif ch == curses.ascii.ENQ:                           # ^e
            self.pos = len(self.line)
        elif ch in (curses.ascii.ACK, curses.KEY_RIGHT):
            if self.pos < len(self.line):
                self.pos +=1
        elif ch == curses.ascii.VT:                            # ^k
            self.line = self.line[:self.pos]
        else:
            return curses.textpad.Textbox.do_command(self, ch)
        self.show_text()
        return 1

    def edit(self):
        curses.doupdate()
        r = curses.textpad.Textbox.edit(self)
        if self.cancelled:
            return None
        return string.join(self.line, "")


class Grub:
    ENTRY_WIN_LINES = 8
    def __init__(self, file, fs = None):
        self.screen = None
        self.entry_win = None
        self.text_win = None
        if file:
            self.read_config(file, fs)

    def draw_main_windows(self):
        if self.screen is None: #only init stuff once
            self.screen = curses.initscr()
            self.screen.timeout(1000)
            if hasattr(curses, 'use_default_colors'):
                try:
                    curses.use_default_colors()
                except:
                    pass # Not important if we can't use colour
            enable_cursor(False)
            self.entry_win = curses.newwin(Grub.ENTRY_WIN_LINES + 2, 74, 2, 1)
            self.text_win = curses.newwin(10, 70, 12, 5)
            curses.def_prog_mode()

        curses.reset_prog_mode()
        self.screen.erase()

        # create basic grub screen with a box of entries and a textbox
        self.screen.addstr(1, 4, "pyGRUB  version %s" %(PYGRUB_VER,))
        self.entry_win.box()
        self.screen.noutrefresh()

    def fill_entry_list(self):
        self.entry_win.erase()
        self.entry_win.box()

        maxy = self.entry_win.getmaxyx()[0]-3 # maxy - 2 for the frame + index
        if self.selected_image > self.start_image + maxy:
            self.start_image = self.selected_image
        if self.selected_image < self.start_image:
            self.start_image = self.selected_image

        for y in range(self.start_image, len(self.cf.images)):
            i = self.cf.images[y]
            if y > self.start_image + maxy:
                break
            if y == self.selected_image:
                self.entry_win.attron(curses.A_REVERSE)
            self.entry_win.addstr(y + 1 - self.start_image, 2, i.title.expandtabs().ljust(70))
            if y == self.selected_image:
                self.entry_win.attroff(curses.A_REVERSE)
        self.entry_win.noutrefresh()

    def edit_entry(self, origimg):
        def draw():
            self.draw_main_windows()

            self.text_win.addstr(0, 0, "Use the U and D keys to select which entry is highlighted.")
            self.text_win.addstr(1, 0, "Press 'b' to boot, 'e' to edit the selected command in the")
            self.text_win.addstr(2, 0, "boot sequence, 'c' for a command-line, 'o' to open a new line")
            self.text_win.addstr(3, 0, "after ('O' for before) the selected line, 'd' to remove the")
            self.text_win.addstr(4, 0, "selected line, or escape to go back to the main menu.")
            self.text_win.addch(0, 8, curses.ACS_UARROW)
            self.text_win.addch(0, 14, curses.ACS_DARROW)
            (y, x) = self.text_win.getmaxyx()
            self.text_win.move(y - 1, x - 1)
            self.text_win.noutrefresh()

        curline = 0
        pos = 0
        img = copy.deepcopy(origimg)
        while 1:
            draw()
            self.entry_win.erase()

            rs = 0
            re = len(img.lines)
            idp = 1
            if re > Grub.ENTRY_WIN_LINES:
                rs = curline - pos
                re = rs + Grub.ENTRY_WIN_LINES

            for idx in range(rs, re):
                # current line should be highlighted
                if idx == curline:
                    self.entry_win.attron(curses.A_REVERSE)

                # trim the line
                l = img.lines[idx].expandtabs().ljust(70)
                if len(l) > 70:
                    l = l[:69] + ">"

                self.entry_win.addstr(idp, 2, l)
                if idx == curline:
                    self.entry_win.attroff(curses.A_REVERSE)
                idp += 1
            self.entry_win.box()
            self.entry_win.noutrefresh()
            curses.doupdate()

            c = self.screen.getch()
            if c in (ord('q'), 27): # 27 == esc
                break
            elif c == curses.KEY_UP:
                curline -= 1
                if pos > 0:
                    pos -= 1
            elif c == curses.KEY_DOWN:
                curline += 1
                if pos < Grub.ENTRY_WIN_LINES - 1:
                    pos += 1
            elif c == ord('b'):
                self.isdone = True
                break
            elif c == ord('e'):
                l = self.edit_line(img.lines[curline])
                if l is not None:
                    img.set_from_line(l, replace = curline)
            elif c == ord('d'):
                img.lines.pop(curline)
            elif c == ord('o'):
                img.lines.insert(curline+1, "")
                curline += 1
            elif c == ord('O'):
                img.lines.insert(curline, "")
            elif c == ord('c'):
                self.command_line_mode()
                if self.isdone:
                    return

            # bound at the top and bottom
            if curline < 0:
                curline = 0
            elif curline >= len(img.lines):
                curline = len(img.lines) - 1

        if self.isdone:
           origimg.reset(img.lines)

    def edit_line(self, line):
        self.screen.erase()
        self.screen.addstr(1, 2, "[ Minimal BASH-like line editing is supported.  ")
        self.screen.addstr(2, 2, "  ESC at any time cancels.  ENTER at any time accepts your changes. ]")
        self.screen.noutrefresh()

        t = GrubLineEditor(self.screen, 5, 2, line)
        enable_cursor(True)
        ret = t.edit()
        if ret:
            return ret
        return None

    def command_line_mode(self):
        self.screen.erase()
        self.screen.addstr(1, 2, "[ Minimal BASH-like line editing is supported.  ESC at any time ")
        self.screen.addstr(2, 2, "  exits.  Typing 'boot' will boot with your entered commands. ] ")
        self.screen.noutrefresh()

        y = 5
        lines = []
        while 1:
            t = GrubLineEditor(self.screen, y, 2)
            enable_cursor(True)
            ret = t.edit()
            if ret:
                if ret in ("quit", "return"):
                    break
                elif ret != "boot":
                    y += 1
                    lines.append(ret)
                    continue

                # if we got boot, then we want to boot the entered image
                img = self.cf.new_image("entered", lines)
                self.cf.add_image(img)
                self.selected_image = len(self.cf.images) - 1
                self.isdone = True
                break

            # else, we cancelled and should just go back
            break

    def read_config(self, fn, fs = None):
        """Read the given file to parse the config.  If fs = None, then
        we're being given a raw config file rather than a disk image."""

        if not os.access(fn, os.R_OK):
            raise RuntimeError("Unable to access %s" %(fn,))

        cfg_list = [(x,grub.GrubConf.Grub2ConfigFile) for x in ["/boot/grub/grub.cfg", "/grub/grub.cfg",
                        "/boot/grub2/grub.cfg", "/grub2/grub.cfg"]] + \
                   [(x,grub.ExtLinuxConf.ExtLinuxConfigFile) for x in ["/boot/isolinux/isolinux.cfg",
                        "/boot/extlinux/extlinux.conf",
                        "/boot/extlinux.conf",
                        "/extlinux/extlinux.conf",
                        "/extlinux.conf"]] + \
                   [(x,grub.GrubConf.GrubConfigFile) for x in ["/boot/grub/menu.lst", "/boot/grub/grub.conf",
                        "/grub/menu.lst", "/grub/grub.conf"]]

        if not fs:
            # set the config file and parse it
            for f,parser in cfg_list:
                self.cf = parser()
                self.cf.filename = fn
                self.cf.parse()
                return

        for f,parser in cfg_list:
            if fs.file_exists(f):
                print("Using %s to parse %s" % (parser,f), file=sys.stderr)
                self.cf = parser()
                self.cf.filename = f
                break
        if self.__dict__.get('cf', None) is None:
            raise RuntimeError("couldn't find bootloader config file in the image provided.")
        f = fs.open_file(self.cf.filename)
        # limit read size to avoid pathological cases
        buf = f.read(FS_READ_MAX)
        del f
        if sys.version_info[0] < 3:
            self.cf.parse(buf)
        else:
            self.cf.parse(buf.decode())

    def image_index(self):
        if isinstance(self.cf.default, int):
            sel = self.cf.default
        elif self.cf.default.isdigit():
            sel = int(self.cf.default)
        else:
            # We don't fully support submenus. Look for the leaf value in
            # "submenu0>submenu1>...>menuentry" and hope that it's unique.
            title = self.cf.default
            while 1:
                try:
                    title = re.search('(\S)>(\S.+$)',title).group(2)
                except AttributeError:
                    break

            # Map string to index in images array
            sel = 0
            for i in range(len(self.cf.images)):
                if self.cf.images[i].title == title:
                    sel = i
                    break

        # If the selected (default) image doesn't exist we select the first entry
        if sel > len(self.cf.images):
            logging.warning("Default image not found")
            sel = 0

        return sel

    def run(self):
        timeout = int(self.cf.timeout)
        self.selected_image = self.image_index()

        self.isdone = False
        while not self.isdone:
            self.run_main(timeout)
            timeout = -1

        return self.selected_image

    def run_main(self, timeout = -1):
        def draw():
            # set up the screen
            self.draw_main_windows()

            if not self.cf.hasPassword() or self.cf.hasPasswordAccess():
                self.text_win.addstr(0, 0, "Use the U and D keys to select which entry is highlighted.")
                self.text_win.addstr(1, 0, "Press enter to boot the selected OS, 'e' to edit the")
                self.text_win.addstr(2, 0, "commands before booting, 'a' to modify the kernel arguments ")
                self.text_win.addstr(3, 0, "before booting, or 'c' for a command line.")

            else:
                self.text_win.addstr(0, 0, "Use the U and D keys to select which entry is highlighted.")
                self.text_win.addstr(1, 0, "Press enter to boot the selected OS or `p` to enter a")
                self.text_win.addstr(2, 0, "password to unlock the next set of features.")

            self.text_win.addch(0, 8, curses.ACS_UARROW)
            self.text_win.addch(0, 14, curses.ACS_DARROW)
            (y, x) = self.text_win.getmaxyx()
            self.text_win.move(y - 1, x - 1)
            self.text_win.noutrefresh()

        # now loop until we hit the timeout or get a go from the user
        mytime = 0
        self.start_image = 0
        while (timeout == -1 or mytime < int(timeout)):
            draw()
            if timeout != -1 and mytime != -1:
                self.screen.addstr(20, 5, "Will boot selected entry in %2d seconds"
                                   %(int(timeout) - mytime))
            else:
                self.screen.addstr(20, 5, " " * 80)
            self.fill_entry_list()
            curses.doupdate()

            c = self.screen.getch()
            if c == -1:
                # Timed out waiting for a keypress
                if mytime != -1:
                    mytime += 1
                    if mytime >= int(timeout):
                        self.isdone = True
                        break
            else:
                # received a keypress: stop the timer
                mytime = -1
                self.screen.timeout(-1)

            # handle keypresses
            if c == ord('c') and self.cf.hasPasswordAccess():
                self.command_line_mode()
                break
            elif c == ord('a') and self.cf.hasPasswordAccess():
                # find the kernel line, edit it and then boot
                img = self.cf.images[self.selected_image]
                for line in img.lines:
                    if line.startswith("kernel") or line.startswith("linux"):
                        l = self.edit_line(line)
                        if l is not None:
                            img.set_from_line(l, replace = True)
                            self.isdone = True
                            break
                break
            elif c == ord('e') and self.cf.hasPasswordAccess():
                img = self.cf.images[self.selected_image]
                self.edit_entry(img)
                break
            elif c == ord('p') and self.cf.hasPassword():
                self.text_win.addstr(6, 1, "Password: ")
                pwd = self.text_win.getstr(6, 8)
                if not self.cf.checkPassword(pwd):
                    self.text_win.addstr(6, 1, "Password: ")
                    if self.cf.passExc is not None:
                        self.text_win.addstr(7, 0, "Exception: %s"
                                                  % self.cf.passExc)
                    else:
                        self.text_win.addstr(7, 0, "Failed!")
                    self.cf.setPasswordAccess( False )
                else:
                    self.cf.setPasswordAccess( True )
                break
            elif c in (curses.KEY_ENTER, ord('\n'), ord('\r')):
                self.isdone = True
                break
            elif c == curses.KEY_UP:
                self.selected_image -= 1
            elif c == curses.KEY_DOWN:
                self.selected_image += 1
#            elif c in (ord('q'), 27): # 27 == esc
#                self.selected_image = -1
#                self.isdone = True
#                break

            # bound at the top and bottom
            if self.selected_image < 0:
                self.selected_image = 0
            elif self.selected_image >= len(self.cf.images):
                self.selected_image = len(self.cf.images) - 1

def get_entry_idx(cf, entry):
    # first, see if the given entry is numeric
    try:
        idx = int(entry)
        return idx
    except ValueError:
        pass

    # it's not, now check the labels for a match
    for i in range(len(cf.images)):
        if entry == cf.images[i].title:
            return i

    return None

def run_grub(file, entry, fs, cfg_args):
    global g
    global sel

    def run_main(scr, *args):
        global sel
        global g
        sel = g.run()

    g = Grub(file, fs)

    if list_entries:
        for i in range(len(g.cf.images)):
            img = g.cf.images[i]
            print("title: %s" % img.title)
            print("  root: %s" % img.root)
            print("  kernel: %s" % img.kernel[1])
            print("  args: %s" % img.args)
            print("  initrd: %s" % img.initrd[1])

    if interactive and not list_entries:
        curses.wrapper(run_main)
    else:
        sel = g.image_index()

    # set the entry to boot as requested
    if entry is not None:
        idx = get_entry_idx(g.cf, entry)
        if idx is not None and idx >= 0 and idx < len(g.cf.images):
           sel = idx

    if sel == -1:
        print("No kernel image selected!")
        sys.exit(1)

    try:
        img = g.cf.images[sel]
    except IndexError:
        img = g.cf.images[0]

    grubcfg = { "kernel": None, "ramdisk": None, "args": "" }

    grubcfg["kernel"] = img.kernel[1]
    if img.initrd:
        grubcfg["ramdisk"] = img.initrd[1]
    if img.args:
        grubcfg["args"] += img.args
    if cfg_args:
        grubcfg["args"] += " " + cfg_args

    return grubcfg

# If nothing has been specified, look for a Solaris domU. If found, perform the
# necessary tweaks.
def sniff_solaris(fs, cfg):
    if not fs.file_exists("/platform/i86xpv/kernel/unix") and \
       not fs.file_exists("/platform/i86xpv/kernel/amd64/unix"):
        return cfg

    if not cfg["kernel"]:
        if fs.file_exists("/platform/i86xpv/kernel/amd64/unix"):
            cfg["kernel"] = "/platform/i86xpv/kernel/amd64/unix"
            cfg["ramdisk"] = "/platform/i86pc/amd64/boot_archive"
        elif fs.file_exists("/platform/i86xpv/kernel/unix"):
            cfg["kernel"] = "/platform/i86xpv/kernel/unix"
            cfg["ramdisk"] = "/platform/i86pc/boot_archive"
        else:
            return cfg

    # Unpleasant. Typically we'll have 'root=foo -k' or 'root=foo /kernel -k',
    # and we need to maintain Xen properties (root= and ip=) and the kernel
    # before any user args.

    xenargs = ""
    userargs = ""

    if not cfg["args"]:
        cfg["args"] = cfg["kernel"]
    else:
        for arg in cfg["args"].split():
            if re.match("^root=", arg) or re.match("^ip=", arg):
                xenargs += arg + " "
            elif arg != cfg["kernel"]:
                userargs += arg + " "
        cfg["args"] = xenargs + " " + cfg["kernel"] + " " + userargs

    return cfg

def sniff_netware(fs, cfg):
    if not fs.file_exists("/nwserver/xnloader.sys"):
        return cfg

    if not cfg["kernel"]:
        cfg["kernel"] = "/nwserver/xnloader.sys"

    return cfg

def format_sxp(kernel, ramdisk, args):
    s = "linux (kernel %s)" % repr(kernel)
    if ramdisk:
        s += "(ramdisk %s)" % repr(ramdisk)
    if args:
        s += "(args %s)" % repr(args)
    return s

def format_simple(kernel, ramdisk, args, sep):
    for check in (kernel, ramdisk, args):
        if check is not None and sep in check:
            raise RuntimeError("simple format cannot represent delimiter-containing value")
    s = ("kernel %s" % kernel) + sep
    if ramdisk:
        s += ("ramdisk %s" % ramdisk) + sep
    if args:
        s += ("args %s" % args) + sep
    s += sep
    return s

if __name__ == "__main__":
    sel = None

    def usage():
        print("Usage: %s [-q|--quiet] [-i|--interactive] [-l|--list-entries] [-n|--not-really] [--output=] [--kernel=] [--ramdisk=] [--args=] [--entry=] [--output-directory=] [--output-format=sxp|simple|simple0] [--runas=] [--offset=] <image>" %(sys.argv[0],), file=sys.stderr)

    def copy_from_image(fs, file_to_read, file_type, fd_dst, path_dst, not_really):
        if not_really:
            if fs.file_exists(file_to_read):
                return "<%s:%s>" % (file_type, file_to_read)
            else:
                sys.exit("The requested %s file does not exist" % file_type)
        try:
            datafile = fs.open_file(file_to_read)
        except Exception as e:
            print(e, file=sys.stderr)
            sys.exit("Error opening %s in guest" % file_to_read)
        dataoff = 0
        while True:
            data = datafile.read(FS_READ_MAX, dataoff)
            if len(data) == 0:
                os.close(fd_dst)
                del datafile
                return
            try:
                os.write(fd_dst, data)
            except Exception as e:
                print(e, file=sys.stderr)
                if path_dst:
                    os.unlink(path_dst)
                del datafile
                sys.exit("Error writing temporary copy of "+file_type)
            dataoff += len(data)

    try:
        opts, args = getopt.gnu_getopt(sys.argv[1:], 'qilnh::',
                                   ["quiet", "interactive", "list-entries", "not-really", "help",
                                    "output=", "output-format=", "output-directory=", "offset=",
                                    "runas=", "entry=", "kernel=",
                                    "ramdisk=", "args=", "isconfig", "debug"])
    except getopt.GetoptError:
        usage()
        sys.exit(1)

    if len(args) < 1:
        usage()
        sys.exit(1)
    file = args[0]
    fs = None
    output = None
    entry = None
    interactive = True
    list_entries = False
    isconfig = False
    part_offs = None
    debug = False
    not_really = False
    output_format = "sxp"
    output_directory = "/var/run/xen/pygrub/"
    uid = None

    # what was passed in
    incfg = { "kernel": None, "ramdisk": None, "args": "" }
    # what grub or sniffing chose
    chosencfg = { "kernel": None, "ramdisk": None, "args": None }
    # what to boot
    bootcfg = { "kernel": None, "ramdisk": None, "args": None }

    for o, a in opts:
        if o in ("-q", "--quiet"):
            interactive = False
        elif o in ("-i", "--interactive"):
            interactive = True
        elif o in ("-l", "--list-entries"):
            list_entries = True
        elif o in ("-n", "--not-really"):
            not_really = True
        elif o in ("-h", "--help"):
            usage()
            sys.exit()
        elif o in ("--output",):
            if a != "-":
                output = a
        elif o in ("--runas",):
            try:
                uid = int(a)
            except ValueError:
                print("runas value must be an integer user id")
                usage()
                sys.exit(1)
        elif o in ("--kernel",):
            incfg["kernel"] = a
        elif o in ("--ramdisk",):
            incfg["ramdisk"] = a
        elif o in ("--args",):
            incfg["args"] = a
        elif o in ("--offset",):
            try:
                part_offs = [ int(a) ]
            except ValueError:
                print("offset value must be an integer")
                usage()
                sys.exit(1)
        elif o in ("--entry",):
            entry = a
            # specifying the entry to boot implies non-interactive
            interactive = False
        elif o in ("--isconfig",):
            isconfig = True
        elif o in ("--debug",):
            debug = True
        elif o in ("--output-format",):
            if a not in ["sxp", "simple", "simple0"]:
                print("unknown output format %s" % a)
                usage()
                sys.exit(1)
            output_format = a
        elif o in ("--output-directory",):
            if not os.path.isdir(a):
                print("%s is not an existing directory" % a)
                sys.exit(1)
            output_directory = a + '/'

    if debug:
        logging.basicConfig(level=logging.DEBUG)

    if interactive and uid:
        print("In order to use --runas, you must also set --entry or -q", file=sys.stderr)
        sys.exit(1)

    try:
        os.makedirs(output_directory, 0o700)
    except OSError as e:
        if (e.errno == errno.EEXIST) and os.path.isdir(output_directory):
            pass
        else:
            raise

    if not_really:
        fd_kernel =  path_kernel = fd_ramdisk = path_ramdisk = None
    else:
        (fd_kernel, path_kernel) = tempfile.mkstemp(prefix="boot_kernel.",
                                                    dir=output_directory)
        (fd_ramdisk, path_ramdisk) = tempfile.mkstemp(prefix="boot_ramdisk.",
                                                      dir=output_directory)

    if output is None:
        fd = sys.stdout.fileno()
    else:
        fd = os.open(output, os.O_WRONLY)

    if uid:
        file = depriv(output_directory, output, file, uid, path_kernel, path_ramdisk)

    # debug
    if isconfig:
        chosencfg = run_grub(file, entry, fs, incfg["args"])
        print("  kernel: %s" % chosencfg["kernel"])
        if chosencfg["ramdisk"]:
            print("  initrd: %s" % chosencfg["ramdisk"])
        print("  args: %s" % chosencfg["args"])
        sys.exit(0)

    # if boot filesystem is set then pass to fsimage.open
    bootfsargs = '"%s"' % incfg["args"]
    bootfsgroup = re.findall('zfs-bootfs=(.*?)[\s\,\"]', bootfsargs)
    if bootfsgroup:
        bootfsoptions = bootfsgroup[0]
    else:
        bootfsoptions = ""

    # get list of offsets into file which start partitions
    if part_offs is None:
        part_offs = get_partition_offsets(file)

    for offset in part_offs:
        try:
            fs = xenfsimage.open(file, offset, bootfsoptions)

            chosencfg = sniff_solaris(fs, incfg)

            if not chosencfg["kernel"]:
                chosencfg = sniff_netware(fs, incfg)

            if not chosencfg["kernel"]:
                chosencfg = run_grub(file, entry, fs, incfg["args"])

            # Break as soon as we've found the kernel so that we continue
            # to use this fsimage object
            if chosencfg["kernel"]:
                break
            fs = None

        except:
            # IOErrors raised by fsimage.open
            # RuntimeErrors raised by run_grub if no menu.lst present
            if debug:
                traceback.print_exc()
            fs = None
            continue

    if list_entries:
        sys.exit(0)

    # Did looping through partitions find us a kernel?
    if fs is None:
        raise RuntimeError("Unable to find partition containing kernel")

    copy_from_image(fs, chosencfg["kernel"], "kernel",
                    fd_kernel, None if uid else path_kernel, not_really)
    bootcfg["kernel"] = path_kernel

    if chosencfg["ramdisk"]:
        try:
            copy_from_image(fs, chosencfg["ramdisk"], "ramdisk",
                            fd_ramdisk, None if uid else path_ramdisk, not_really)
        except:
            if not uid and not not_really:
                    os.unlink(path_kernel)
            raise
        bootcfg["ramdisk"] = path_ramdisk
    else:
        initrd = None
        if not uid and not not_really:
            os.unlink(path_ramdisk)

    args = None
    if chosencfg["args"]:
        zfsinfo = xenfsimage.getbootstring(fs)
        if zfsinfo is not None:
            e = re.compile("zfs-bootfs=[\w\-\.\:@/]+" )
            (chosencfg["args"],count) = e.subn(zfsinfo, chosencfg["args"])
            if count == 0:
               chosencfg["args"] += " -B %s" % zfsinfo
        args = chosencfg["args"]

    if output_format == "sxp":
        ostring = format_sxp(bootcfg["kernel"], bootcfg["ramdisk"], args)
    elif output_format == "simple":
        ostring = format_simple(bootcfg["kernel"], bootcfg["ramdisk"], args, "\n")
    elif output_format == "simple0":
        ostring = format_simple(bootcfg["kernel"], bootcfg["ramdisk"], args, "\0")

    sys.stdout.flush()
    if sys.version_info[0] < 3:
        os.write(fd, ostring)
    else:
        os.write(fd, ostring.encode())

